from dna.Dialogs import *
from dna.SiteFrequencySpectrum import *
from dna.Parameters import *
from dna.PolymorphismTables import PolymorphismTables
import dna.SiteFrequencySpectrum as sfs
from dna.PolyspectSpace import *
import matplotlib.pyplot as plt
import infiniteSitesTools as ist
from dna.R2_Tools import R2

def main():
	
	modelPop = Dialog_MEMES()
	memes = modelPop.getResponse()
	
	diaPop = Dialog_Population()
	pop = diaPop.getResponse()

	nullDemography = Demography( pop.N )
	nullDemography.addEpoch( 10000, 0, True )
	print nullDemography
	print pop.getMSString_Null(), nullDemography.getMSString()
	
	demDialog = None
	
	if memes == modelPop.memes[0]:
		demDialog = Dialog_Demography_C( pop.N )
	elif memes == modelPop.memes[1]:
		demDialog = Dialog_Demography_G( pop.N )
	elif memes == modelPop.memes[2]:
		demDialog = Dialog_Demography_CC( pop.N )
	elif memes == modelPop.memes[3]:
		demDialog = Dialog_Demography_CG( pop.N )
	else:
		return
		
	altDemography =  demDialog.getResponse()
		
	simParams = SimulationParameters( pop, nullDemography )
	print simParams

	pt = PolymorphismTables()
	msCommand = simParams.getMSString_Null()
	msOut = pt.runMS(msCommand)
	pt.readMS(None, msOut)

	mySFSs = SiteFrequencySpectrum(pt)
	
	pt = PolymorphismTables()
	msCommand = simParams.getMSString_Alternative(altDemography )
	msOut = pt.runMS(msCommand)
	pt.readMS(None, msOut)
	myAltSFSs = SiteFrequencySpectrum(pt)
	
	mySpace = PolyspectSpace( simParams, altDemography, mySFSs, myAltSFSs)
	
	point = mySpace.getCurrentPoint()
	print point
	
	for i in range(0,100):
		newPoint = mySpace.getNextPoint()
		if newPoint == None:
			print "\nFound peak for now"
			break
			
		print newPoint
	
	rr = R2( simParams )
	r2power = rr.getPower(altDemography)
	print 'R2 power:\t' + str(r2power)

	makePlot( mySpace, r2power )


def makePlot( mySpace, r2power ):
	tjd = ist.normalizeToFirst( mySpace.myPowerPath[0].myPolyspect.myOperator )
	flf = ist.normalizeToFirst( mySpace.myPowerPath[3].myPolyspect.myOperator )
	dmd = ist.normalizeToFirst( mySpace.myPowerPath[-1].myPolyspect.myOperator )


	myPowers = []
	myPowers.append( mySpace.myPowerPath[0].myPower)
	myPowers.append( mySpace.myPowerPath[2].myPower)
	myPowers.append( mySpace.myPowerPath[3].myPower)
	myPowers.append( r2power )
	myPowers.append( mySpace.myPowerPath[-1].myPower)
	
	fig = plt.figure()
	ind = np.arange(30)
	
	ax = fig.add_subplot(2,2,1)
	ax.bar( ind, tjd, color='m' )
	ax.set_title('Tajimas D')

	ax = fig.add_subplot(2,2,2)
	ax.bar( ind, flf, color='b')
	ax.set_title('Fu Lis F')

	ax = fig.add_subplot(2,2,3)
	ax.bar( ind, dmd, color='y')
	ax.set_title('Davids M')

	ind = np.arange(5)
	ax = fig.add_subplot(2,2,4)
	rects = ax.bar( ind, myPowers, color='r')
	ax.set_title('Powers')
	ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
	ax.set_xticks(ind+0.35)
	ax.set_xticklabels( ('TajD', 'FuLiD', 'FuLiF', 'R2', 'DavidsM') )

	def autolabel(rects):
		# attach some text labels
		for rect in rects:
			height = rect.get_height()
			ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(100*height), ha='center', va='bottom')

	autolabel(rects)

	print "\n*** Final path ***"
	print mySpace
	print "\nProgram complete"
	
	mySpace.getSlopeAtMaxPoint()

	plt.show()


	#r = codebox( "*** Final path ***", "Results", str(mySpace) )
	
main()
